icu_benchmarks/data/loader.py (2)

115-134: ⚠️ Potential issue | 🔴 Critical

Critical: in-place mutation of cached _label_arrays[stay_id] corrupts pad_mask on subsequent calls.

On Line 116, labels is a live reference to the array stored in self._label_arrays[stay_id]. When len(labels) > 1 and length_diff <= 0 (i.e., a stay whose length equals maxlen), neither the len(labels) == 1 nor the length_diff > 0 branch allocates a new array, so labels[not_labeled] = -1 on Line 133 mutates the cached entry in place.

After the first call, the cached labels have -1 where NaNs used to be. On any subsequent rebuild (e.g., when ram_cache=False, or if ram_cache(True) is ever called again), np.argwhere(np.isnan(labels)) returns an empty array, so the pad_mask[not_labeled] = 0 line on Line 134 is skipped and pad_mask is returned as all ones at positions that should have been masked out — silently feeding unlabeled timesteps into the loss.

window is never mutated, so only labels needs a defensive copy.

🐛 Proposed fix

         window = self._feat_arrays[stay_id]
-        labels = self._label_arrays[stay_id]
+        # Copy to avoid mutating the cached entry below (labels[not_labeled] = -1).
+        labels = self._label_arrays[stay_id].copy()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@icu_benchmarks/data/loader.py` around lines 115 - 134, The issue is in-place
mutation of the cached labels array (self._label_arrays[stay_id]) causing
persistent -1s that break pad_mask later; fix by making a defensive copy of the
labels before any mutation (e.g., assign labels =
self._label_arrays[stay_id].copy() or np.array(self._label_arrays[stay_id],
copy=True) immediately after retrieving it), and ensure any branches that build
new label arrays (the len(labels) == 1 concatenation and the length_diff > 0
padding branch) operate on and return this copied/allocated array so subsequent
writes like labels[not_labeled] = -1 do not mutate the cached self._label_arrays
entry (refer to variables stay_id, labels, self._label_arrays, not_labeled,
pad_mask).

---

174-195: ⚠️ Potential issue | 🔴 Critical

Critical: float32 cast of row_indicators corrupts stay-id precision in exported predictions.

row_indicators carries the GROUP (stay_id) and, for temporal data, SEQUENCE (time in hours). These are actively passed to _save_model_outputs() and exported to pred_indicators.csv for downstream analysis. float32 can only represent integers exactly up to 2²⁴ ≈ 16.7 million; hospital-derived stay ids (e.g., MIMIC-IV) routinely exceed this. Unconditionally casting to float32 on line 194 silently collapses distinct ids, corrupting the CSV export and breaking any attempt to match predictions back to original stays.

This is a regression: commit 52117f4 introduced conditional mps-based casting (if self.mps: ...else: native dtype), but the self.mps flag (line 53) is set but never used in to_tensor(). The condition was removed and replaced with unconditional float32 casting.

Only data and labels require float32 for tensor operations; row_indicators is metadata and should retain its native dtype.

Fix for to_tensor

     def to_tensor(self) -> tuple[Tensor, Tensor, Tensor]:
         data, labels, row_indicators = self.get_data_and_labels()
-        # Always use float32 for memory efficiency and MPS compatibility
         return (
             from_numpy(data),
             from_numpy(labels),
-            from_numpy(row_indicators.astype(np.float32)),
+            from_numpy(row_indicators),
         )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@icu_benchmarks/data/loader.py` around lines 174 - 195, The to_tensor method
currently forces row_indicators to float32 which corrupts stay_id precision;
update to_tensor (and/or get_data_and_labels usage) so only data and labels are
cast to float32 before creating tensors and row_indicators is passed through
with its native integer dtype (do not call .astype(np.float32) on
row_indicators). Use the existing self.mps flag if you need MPS-specific casting
for data/labels, but ensure row_indicators stays as the original numpy dtype
when calling from_numpy(row_indicators) so exported pred_indicators.csv keeps
exact stay_id/sequence values.